import json
import os
import Agently

# agent_factory = (
#      Agently.AgentFactory()
#      .set_settings("current_model", "OAIClient")
#      .set_settings("model.OAIClient.url", os.environ["OPENAI_API_BASE_ZHIHU"])
#      .set_settings("model.OAIClient.auth", { "api_key": os.getenv("OPENAI_API_KEY_ZHIHU")})
#      .set_settings("model.OAIClient.options", { "model": "gpt-3.5-turbo" })
# )
agent_factory = (
        Agently.AgentFactory()
        # 给Agent工厂实例提供设置项：
        ## 将默认模型系列设置为智谱系列
        .set_settings("current_model", "ZhipuAI")
        ## 提供你的智谱线上模型API-KEY
        .set_settings("model.ZhipuAI.auth", { "api_key": os.environ["ZHIPU_API_KEY"] })
        ## 模型默认为GLM-4，如果想要指定为GLM-3-turbo可以解除下面代码的注释
        .set_settings("model.ZhipuAI.options", { "model": "glm-3-turbo" })
)

workflow = Agently.Workflow()

@workflow.chunk()
def initial_translation(inputs,storage):
    """
    Translate the entire text as one chunk using an LLM.

    Args:
        State : 状态对象

    Returns:
        str: The translated text.
    """

    """
        source_lang (str): The source language of the text.
        target_lang (str): The target language for translation.
        source_text (str): The text to be translated.
    """
    source_lang = storage.get("source_lang")
    target_lang = storage.get("target_lang")
    source_text = storage.get("source_text")

    translate_agent = agent_factory.create_agent()
    translate_agent.set_agent_prompt("role",f"You are an expert linguist, specializing in translation from {source_lang} to {target_lang}.")

    translation_prompt = f"""This is an {source_lang} to {target_lang} translation, please provide the {target_lang} translation for this text. \
Do not provide any explanations or text apart from the translation.
{source_lang}: {source_text}

{target_lang}:"""

    translation_1 = (translate_agent.input(translation_prompt).start())
    print("初次翻译结果：\n",translation_1)

    storage.set("translation_1",translation_1)
    storage.set("translate_agent",translate_agent)

    return {"stage":"initial translation","result":translation_1}


@workflow.chunk()
def reflect_on_translation(inputs,storage):
    """
    Use an LLM to reflect on the translation, treating the entire text as one chunk.

    Args:
        state

    Returns:
        str: The LLM's reflection on the translation, providing constructive criticism and suggestions for improvement.
    """

    """
        source_lang (str): The source language of the text.
        target_lang (str): The target language of the translation.
        source_text (str): The original text in the source language.
        translation_1 (str): The initial translation of the source text.
        country (str): Country specified for target language.
    """
    source_lang = storage.get("source_lang")
    target_lang =  storage.get("target_lang")
    source_text = storage.get("source_text")
    translation_1= storage.get("translation_1")
    country: str= storage.get("country")
    replaceText = None

    translate_agent = agent_factory.create_agent()
    if country :
        replaceText = "The final style and tone of the translation should match the style of {target_lang} colloquially spoken in {country}."
    else:
        replaceText = ""

    translate_agent.set_agent_prompt("role",f"You are an expert linguist specializing in translation from {source_lang} to {target_lang}. \
You will be provided with a source text and its translation and your goal is to improve the translation.")

    reflection_prompt = f"""Your task is to carefully read a source text and a translation from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions to improve the translation. \
    {replaceText}

    The source text and initial translation, delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT> and <TRANSLATION></TRANSLATION>, are as follows:

    <SOURCE_TEXT>
    {source_text}
    </SOURCE_TEXT>

    <TRANSLATION>
    {translation_1}
    </TRANSLATION>

    When writing suggestions, pay attention to whether there are ways to improve the translation's \n\
    (i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),\n\
    (ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules, and ensuring there are no unnecessary repetitions),\n\
    (iii) style (by ensuring the translations reflect the style of the source text and takes into account any cultural context),\n\
    (iv) terminology (by ensuring terminology use is consistent and reflects the source text domain; and by only ensuring you use equivalent idioms {target_lang}).\n\

    Write a list of specific, helpful and constructive suggestions for improving the translation.
    Each suggestion should address one specific part of the translation.
    Output only the suggestions and nothing else."""

    reflection = (translate_agent.input(reflection_prompt).start())

    storage.set("reflection",reflection)

    print("改进意见：\n",reflection)
    return {"stage":"reflect on translation","result":reflection}


@workflow.chunk()
def improve_translation(inputs,storage):
    """
    Use the reflection to improve the translation, treating the entire text as one chunk.

    Args:
        state

    Returns:
        str: The improved translation based on the expert suggestions.
    """

    """
        source_lang (str): The source language of the text.
        target_lang (str): The target language for the translation.
        source_text (str): The original text in the source language.
        translation_1 (str): The initial translation of the source text.
        reflection (str): Expert suggestions and constructive criticism for improving the translation.
    """
    source_lang = storage.get("source_lang")
    target_lang = storage.get("target_lang")
    source_text = storage.get("source_text")
    translation_1 = storage.get("translation_1")
    reflection = storage.get("reflection")

    translate_agent = storage.get("translate_agent")
    translate_agent.set_agent_prompt("role", f"You are an expert linguist, specializing in translation editing from {source_lang} to {target_lang}.")

    prompt = f"""Your task is to carefully read, then edit, a translation from {source_lang} to {target_lang}, taking into
account a list of expert suggestions and constructive criticisms.

The source text, the initial translation, and the expert linguist suggestions are delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT>, <TRANSLATION></TRANSLATION> and <EXPERT_SUGGESTIONS></EXPERT_SUGGESTIONS> \
as follows:

<SOURCE_TEXT>
{source_text}
</SOURCE_TEXT>

<TRANSLATION>
{translation_1}
</TRANSLATION>

<EXPERT_SUGGESTIONS>
{reflection}
</EXPERT_SUGGESTIONS>

Please take into account the expert suggestions when editing the translation. Edit the translation by ensuring:

(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),
(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules and ensuring there are no unnecessary repetitions), \
(iii) style (by ensuring the translations reflect the style of the source text)
(iv) terminology (inappropriate for context, inconsistent use), or
(v) other errors.

Output only the new translation and nothing else."""

    translation_2 = translate_agent.input(prompt).start()

    storage.set("translation_2",translation_2)
    print("改进后的翻译：\n", translation_2)

    return {"stage":"improve translation","result":translation_2}


@workflow.chunk_class()
def output_stage_result(inputs,storage):
    print(f"[{inputs['default']['stage']}]:\n", inputs["default"]["result"])
    return

def translate_text(source_lang: str, target_lang: str, source_text: str, country: str = ""):
    """
    Translate a single chunk of text from the source language to the target language.

    This function performs a two-step translation process:
    1. Get an initial translation of the source text.
    2. Reflect on the initial translation and generate an improved translation.

    Returns:
        str: The improved translation of the source text.
    """

    (workflow.connect_to("initial_translation").connect_to("reflect_on_translation").connect_to("improve_translation").connect_to("end"))

    (workflow.chunks["initial_translation"].connect_to("@output_stage_result").connect_to("reflect_on_translation.wait"))
    (workflow.chunks["reflect_on_translation"].connect_to("@output_stage_result").connect_to("improve_translation.wait"))
    (workflow.chunks["improve_translation"].connect_to("@output_stage_result"))

    result = workflow.start(storage = {
        "source_lang": source_lang,
        "target_lang": target_lang,
        "source_text": source_text
    })

    print(json.dumps(result, indent=4, ensure_ascii=False))

if __name__ == "__main__":
    translate_text("English","中文","""Translation Agent: Agentic translation using reflection workflow
    This is a Python demonstration of a reflection agentic workflow for machine translation. The main steps are:

    Prompt an LLM to translate a text from source_language to target_language;
    Have the LLM reflect on the translation to come up with constructive suggestions for improving it;
    Use the suggestions to improve the translation.
    Customizability
    By using an LLM as the heart of the translation engine, this system is highly steerable. For example, by changing the prompts, it is easier using this workflow than a traditional machine translation (MT) system to:

    Modify the output's style, such as formal/informal.
    Specify how to handle idioms and special terms like names, technical terms, and acronyms. For example, including a glossary in the prompt lets you make sure particular terms (such as open source, H100 or GPU) are translated consistently.
    Specify specific regional use of the language, or specific dialects, to serve a target audience. For example, Spanish spoken in Latin America is different from Spanish spoken in Spain; French spoken in Canada is different from how it is spoken in France.
    This is not mature software, and is the result of Andrew playing around with translations on weekends the past few months, plus collaborators (Joaquin Dominguez, Nedelina Teneva, John Santerre) helping refactor the code.

    According to our evaluations using BLEU score on traditional translation datasets, this workflow is sometimes competitive with, but also sometimes worse than, leading commercial offerings. However, we’ve also occasionally gotten fantastic results (superior to commercial offerings) with this approach. We think this is just a starting point for agentic translations, and that this is a promising direction for translation, with significant headroom for further improvement, which is why we’re releasing this demonstration to encourage more discussion, experimentation, research and open-source contributions.

    If agentic translations can generate better results than traditional architectures (such as an end-to-end transformer that inputs a text and directly outputs a translation) -- which are often faster/cheaper to run than our approach here -- this also provides a mechanism to automatically generate training data (parallel text corpora) that can be used to further train and improve traditional algorithms. (See also this article in The Batch on using LLMs to generate training data.)

    Comments and suggestions for how to improve this are very welcome!""")

